{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [基于序列信息的microRNA和gene关系对预测](https://datafountain.cn/competitions/534)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from models import BaseInfo, BaseSeq2Vec, RelationShipModel\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import sys\n",
    "import gc\n",
    "# GCNFrame依赖于torch1.7.*\n",
    "PRETRAINED_PATH = './user_data/'\n",
    "OUTPUT_PATH = './user_data/'\n",
    "DATA_PATH = './datasets/'\n",
    "print(sys.argv)\n",
    "if len(sys.argv) == 4:\n",
    "    DATA_PATH = sys.argv[1]\n",
    "    OUTPUT_PATH = sys.argv[2]\n",
    "    PRETRAINED_PATH = sys.argv[3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 进行简单的预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = pd.read_csv(f'{DATA_PATH}train_dataset/Train.csv')\n",
    "mirna_seqdf = pd.read_csv(f'{DATA_PATH}train_dataset/mirna_seq.csv', index_col=0)\n",
    "gene_seqdf = pd.read_csv(f'{DATA_PATH}train_dataset/gene_seq.csv', index_col=0)\n",
    "predict_set = pd.read_csv(f'{DATA_PATH}test_dataset.csv')# gene,miRNA\n",
    "\n",
    "# 清除换行符\n",
    "mirna_seqdf['seq'] = mirna_seqdf['seq'].apply(lambda x:x.strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 加载已经训练好的特征提取器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(f'{OUTPUT_PATH}base_info.m'):\n",
    "    base_info = BaseInfo.load(f'{OUTPUT_PATH}base_info.m')\n",
    "else:\n",
    "    base_info = BaseInfo()\n",
    "\n",
    "if os.path.exists(f'{PRETRAINED_PATH}rna2vec.m'):\n",
    "    rna2vec = BaseSeq2Vec.load(f'{PRETRAINED_PATH}rna2vec.m').eval()\n",
    "else:\n",
    "    rna2vec = BaseSeq2Vec(base_info.get_RNAf().shape[-1], 32, base_info.get_rna_s2i(), base_info.get_RNAf(), base_info.k).eval()\n",
    "\n",
    "if os.path.exists(f'{PRETRAINED_PATH}dna2vec.m'):\n",
    "    dna2vec = BaseSeq2Vec.load(f'{PRETRAINED_PATH}dna2vec.m').eval()\n",
    "else:\n",
    "    dna2vec = BaseSeq2Vec(base_info.get_DNAf().shape[-1], 32, base_info.get_dna_s2i(), base_info.get_DNAf(), base_info.k).eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 提取训练集特征\n",
    "参照`Train.csv`中DNA与RNA的对应关系，将序列名字换成特征向量，最后得到两个可以相互对应的特征矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def base2feature(base_name:list, seq:pd.Series, transform:BaseSeq2Vec, gaps:list=[0]) -> torch.Tensor:\n",
    "    '''\n",
    "    base_name:碱基序列名字的列表，输出按照顺序组织特征向量\n",
    "    seq:碱基序列名字与对应的序列(名字作为index)\n",
    "    transform:可以将序列转换成特征向量的转换器\n",
    "    '''\n",
    "    if not transform.trained:\n",
    "        print('Warning: Using untrained BaseSeq2Vec')\n",
    "    # 获取需要转换的序列名字的列表\n",
    "    name_list = list(set(base_name))\n",
    "    # 获取其对应的特征\n",
    "    name2vec = dict()\n",
    "    with torch.no_grad():\n",
    "        for name in tqdm(name_list):\n",
    "            tmp = []\n",
    "            for gap in gaps:\n",
    "                tmp.append(\n",
    "                    transform(seq[name], gap).unsqueeze(0)\n",
    "                )\n",
    "            name2vec[name] = torch.concat(tmp, dim=0).unsqueeze(0)\n",
    "    # 按照base_name的顺序组织输出\n",
    "    ans = []\n",
    "    for name in tqdm(base_name):\n",
    "        ans.append(name2vec[name])\n",
    "    return torch.concat(ans, dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取训练集的RNA特征矩阵\n",
    "print('preparing train_rna_feature')\n",
    "if not os.path.exists(f'{OUTPUT_PATH}train_rna_feature.tsr'):\n",
    "    train_rna_feature = base2feature(dataset['miRNA'].tolist(), mirna_seqdf['seq'], rna2vec, [0, 1, 2])\n",
    "    torch.save(train_rna_feature, f'{OUTPUT_PATH}train_rna_feature.tsr')\n",
    "gc.collect()\n",
    "\n",
    "# 获取预测集的RNA特征矩阵\n",
    "print('preparing test_rna_feature')\n",
    "if not os.path.exists(f'{OUTPUT_PATH}test_rna_feature.tsr'):\n",
    "    test_rna_feature = base2feature(predict_set['miRNA'].tolist(), mirna_seqdf['seq'], rna2vec, [0, 1, 2])\n",
    "    torch.save(test_rna_feature, f'{OUTPUT_PATH}test_rna_feature.tsr')\n",
    "gc.collect()\n",
    "\n",
    "# 获取训练集的DNA特征矩阵\n",
    "print('preparing train_dna_feature')\n",
    "if not os.path.exists(f'{OUTPUT_PATH}train_dna_feature.tsr'):\n",
    "    train_dna_feature = base2feature(dataset['gene'].tolist(), gene_seqdf['sequence'], dna2vec, [0, 1, 2])\n",
    "    torch.save(train_dna_feature, f'{OUTPUT_PATH}train_dna_feature.tsr')\n",
    "gc.collect()\n",
    "\n",
    "# 获取预测集的DNA特征矩阵\n",
    "print('preparing test_dna_feature')\n",
    "if not os.path.exists(f'{OUTPUT_PATH}test_dna_feature.tsr'):\n",
    "    test_dna_feature = base2feature(predict_set['gene'].tolist(), gene_seqdf['sequence'], dna2vec, [0, 1, 2])\n",
    "    torch.save(test_dna_feature, f'{OUTPUT_PATH}test_dna_feature.tsr')\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 获取训练集标签"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('preparing labels')\n",
    "if not os.path.exists(f'{OUTPUT_PATH}labels.tsr'):\n",
    "    label_dict = {\n",
    "        'Functional MTI': 1.,\n",
    "        'Non-Functional MTI': 0.}\n",
    "    labels = dataset['label'].map(lambda x: label_dict[x]).values\n",
    "    labels = torch.tensor(labels, dtype=torch.float32)\n",
    "    torch.save(labels, f'{OUTPUT_PATH}labels.tsr')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 搭建用于预测的模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(f'{OUTPUT_PATH}model.pth'):\n",
    "    train_dna_feature = torch.load(f'{OUTPUT_PATH}train_dna_feature.tsr')\n",
    "    train_rna_feature = torch.load(f'{OUTPUT_PATH}train_rna_feature.tsr')\n",
    "    labels = torch.load(f'{OUTPUT_PATH}labels.tsr')\n",
    "    \n",
    "    model = RelationShipModel(train_rna_feature.shape[-1], train_rna_feature.shape[-2])\n",
    "    model.fit(train_dna_feature, train_rna_feature, labels, 20, 150, 1e-3)\n",
    "    torch.save(model, f'{OUTPUT_PATH}model.pth')\n",
    "    \n",
    "    del train_dna_feature, train_rna_feature, labels\n",
    "    gc.collect()\n",
    "model = torch.load(f'{OUTPUT_PATH}model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 开始预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dna_feature = torch.load(f'{OUTPUT_PATH}test_dna_feature.tsr')\n",
    "test_rna_feature = torch.load(f'{OUTPUT_PATH}test_rna_feature.tsr')\n",
    "\n",
    "ans = model.predict(test_dna_feature, test_rna_feature, 150)\n",
    "predict_set['ans'] = ans\n",
    "predict_set['results'] = predict_set['ans'].map(lambda x:int(x+0.5))\n",
    "predict_set.to_csv(f'{OUTPUT_PATH}ans.csv', sep=',', index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AI",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
